import open3d as o3d
import numpy as np
pcd = o3d.io.read_point_cloud("../datasets/standford_cloud_data/China dragon.pcd")
pcd.paint_uniform_color([0.5, 0.5, 0.5])#把所有点渲染为灰色
pcd_tree = o3d.geometry.KDTreeFlann(pcd)#建立KD树索引

#K近邻搜索
pcd.colors[10000] = [1, 0, 0]#给定查询点并渲染为红色


[k, idx, _] = pcd_tree.search_knn_vector_3d(pcd.points[10000], 200)#K近邻搜索
np.asarray(pcd.colors)[idx[1:], :] = [0, 1, 0]#K邻域的点，渲染为绿色

#半径搜索
pcd.colors[5000] = [1, 0, 0]#给定查询点并渲染为红色
[k1, idx1, _] = pcd_tree.search_radius_vector_3d(pcd.points[5000], 0.02)#半径搜索
np.asarray(pcd.colors)[idx1[1:], :] = [0, 0, 1]#半径搜索结果并渲染为蓝色

#混合搜索
pcd.colors[30000] = [1, 1, 0]#给定查询点并渲染为黄色
[k2, idx2, _] = pcd_tree.search_hybrid_vector_3d(pcd.points[30000], 0.05,200)#K近邻搜索
np.asarray(pcd.colors)[idx2[1:], :] = [0, 1, 0.8]#半径搜索结果并渲染为青色
o3d.visualization.draw_geometries([pcd])
